# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os
import datetime
import logging

from absl import app, flags


def _iterate_and_make_hard_link(path):
  if not os.path.exists(path):
    return
  for sub_req_dir in sorted(os.listdir(path)):
    for old_file_name in sorted(os.listdir(os.path.join(path, sub_req_dir))):
      exploded = old_file_name.split('.')
      assert len(exploded) >= 5
      needed = exploded[:-4] + exploded[-2:]
      new_file_name = '.'.join(needed)
      os.link(os.path.join(path, sub_req_dir, old_file_name),
              os.path.join(path, sub_req_dir, new_file_name))


def _gen_dir_to_process(root, machine, date):
  date_str = date.strftime('%Y%m%d')
  return os.path.join(root, machine, date_str)


def main(argv):
  FLAGS = flags.FLAGS
  if '-' in FLAGS.date:
    from_date_str, to_date_str = FLAGS.date.split('-')
    from_date = datetime.datetime.strptime(from_date_str, '%Y%m%d').date()
    to_date = datetime.datetime.strptime(to_date_str, '%Y%m%d').date()
  else:
    from_date = datetime.datetime.strptime(FLAGS.date, '%Y%m%d').date()
    to_date = from_date

  machine_whitelist = FLAGS.machine.split(',')

  cur_date = from_date
  while cur_date <= to_date:
    for machine in machine_whitelist:
      logging.info('Make hard link for %s %s' % (cur_date, machine))
      _iterate_and_make_hard_link(_gen_dir_to_process(FLAGS.feed_cache_root, machine, cur_date))
    cur_date += datetime.timedelta(days=1)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('feed_cache_root',
                      '/remote/iosg/data/coin/feed_cache',
                      'Input fastfeed cache root dir.')

  flags.DEFINE_string('machine', 'feed-01.ap-northeast-1.aws', 'machine')

  flags.DEFINE_string('date', None, 'Date')

  app.run(main)
